"""
完整的高铁API集成 - 使用12306官方接口
"""
import requests
import json
import re
from datetime import datetime, timedelta
import time

class RailwayAPI:
    def __init__(self):
        self.base_url = "https://kyfw.12306.cn/otn"
        self.session = requests.Session()
        self.station_codes = self._load_station_codes()
        
    def _load_station_codes(self):
        """加载火车站代码映射：优先动态拉取12306站点表，失败则回退内置映射"""
        # 动态拉取 12306 站点表
        remote_urls = [
            "https://kyfw.12306.cn/otn/resources/js/framework/station_name.js",
            "https://kyfw.12306.cn/otn/resources/js/framework/station_name_v10012.js"
        ]
        try:
            for url in remote_urls:
                try:
                    resp = self.session.get(url, timeout=8)
                    if resp.status_code == 200 and resp.text:
                        text = resp.text
                        # station_name.js 格式示例: var station_names = '@bji|北京|BJP|beijing|bj|2|…@shh|上海|SHH|…'
                        import re
                        pattern = re.compile(r"@[^|]*\|([^|]+)\|([A-Z]{3})\|")
                        mapping = {}
                        for m in pattern.finditer(text):
                            zh_name = m.group(1)
                            code = m.group(2)
                            mapping[zh_name] = code
                        if mapping:
                            # 返回全国覆盖映射
                            return mapping
                except Exception:
                    continue
        except Exception:
            pass
        
        # 回退：内置少量映射，保障最小可用
        return {
            '北京': 'BJP', '北京北': 'VAP', '北京东': 'BOP', '北京南': 'VNP', '北京西': 'BXP',
            '上海': 'SHH', '上海南': 'SNH', '上海虹桥': 'AOH', '上海西': 'SXH',
            '广州': 'GZQ', '广州北': 'GBQ', '广州东': 'GGQ', '广州南': 'IZQ', '广州西': 'GXQ',
            '深圳': 'SZQ', '深圳北': 'IOQ', '深圳东': 'BJQ', '深圳西': 'OSQ',
            '杭州': 'HZH', '杭州东': 'HGH', '杭州南': 'XHH',
            '南京': 'NJH', '南京南': 'NKH',
            '武汉': 'WHN', '武汉东': 'LFN',
            '成都': 'CDW', '成都东': 'ICW', '成都南': 'CNW', '成都西': 'CMW',
            '西安': 'XAY', '西安北': 'EAY',
            '重庆': 'CQW', '重庆北': 'CUW', '重庆南': 'CRW', '重庆西': 'CXW',
            '天津': 'TJP', '天津北': 'TBP', '天津南': 'TIP', '天津西': 'TXP',
            '青岛': 'QDK', '青岛北': 'QHK', '青岛西': 'QUK',
            '大连': 'DLT', '大连北': 'DFT',
            '合肥': 'HFH', '合肥南': 'ENH', '合肥西': 'HTH',
            '大连交通大学': 'DLT',
            '合肥南站': 'ENH',
            '北京交通大学': 'BJP',
        }
    
    def get_station_code(self, station_name):
        """获取火车站代码"""
        # 清理站名
        clean_name = station_name.strip()
        for suf in ['站', '火车站', '大学', '学院', '学校', '市', '区', '县', '镇', '乡', '村', '校']:
            clean_name = clean_name.replace(suf, '')
        clean_name = clean_name.strip()
        
        # 直接匹配
        if clean_name in self.station_codes:
            return self.station_codes[clean_name]
        
        # 模糊匹配
        for name, code in self.station_codes.items():
            if clean_name in name or name in clean_name:
                return code
        
        return None
    
    def search_trains(self, from_station, to_station, date=None):
        """搜索火车票"""
        if date is None:
            date = datetime.now().strftime('%Y-%m-%d')
        
        from_code = self.get_station_code(from_station)
        to_code = self.get_station_code(to_station)
        
        if not from_code or not to_code:
            return {
                'success': False,
                'error': f'无法找到车站代码: {from_station} -> {to_station}'
            }
        
        try:
            # 构建查询URL
            url = f"{self.base_url}/leftTicket/query"
            params = {
                'leftTicketDTO.train_date': date,
                'leftTicketDTO.from_station': from_code,
                'leftTicketDTO.to_station': to_code,
                'purpose_codes': 'ADULT'
            }
            
            # 设置请求头
            headers = {
                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
                'Referer': 'https://kyfw.12306.cn/otn/leftTicket/init',
                'Accept': 'application/json, text/javascript, */*; q=0.01',
                'Accept-Language': 'zh-CN,zh;q=0.9,en;q=0.8',
                'X-Requested-With': 'XMLHttpRequest'
            }
            
            # 发送请求
            response = self.session.get(url, params=params, headers=headers, timeout=10)
            
            if response.status_code == 200:
                data = response.json()
                
                if data.get('status') and data.get('data'):
                    trains = self._parse_train_data(data['data']['result'])
                    return {
                        'success': True,
                        'trains': trains,
                        'date': date,
                        'from_station': from_station,
                        'to_station': to_station
                    }
                else:
                    return {
                        'success': False,
                        'error': '暂无车次信息'
                    }
            else:
                return {
                    'success': False,
                    'error': f'请求失败: {response.status_code}'
                }
                
        except Exception as e:
            return {
                'success': False,
                'error': f'查询异常: {str(e)}'
            }
    
    def _parse_train_data(self, train_data):
        """解析火车数据"""
        trains = []
        
        for train_info in train_data:
            if not train_info or '|' not in train_info:
                continue
                
            parts = train_info.split('|')
            
            if len(parts) < 30:
                continue
            
            try:
                train = {
                    'train_number': parts[3],  # 车次
                    'from_station': parts[4],  # 出发站代码
                    'to_station': parts[7],    # 到达站代码
                    'departure_time': parts[8],  # 出发时间
                    'arrival_time': parts[9],    # 到达时间
                    'duration': parts[10],       # 历时
                    'can_buy': parts[11] == 'Y', # 是否可以购买
                    'business_seat': parts[32] if len(parts) > 32 else '',  # 商务座
                    'first_class': parts[31] if len(parts) > 31 else '',    # 一等座
                    'second_class': parts[30] if len(parts) > 30 else '',   # 二等座
                    'soft_sleeper': parts[23] if len(parts) > 23 else '',   # 软卧
                    'hard_sleeper': parts[28] if len(parts) > 28 else '',   # 硬卧
                    'soft_seat': parts[24] if len(parts) > 24 else '',      # 软座
                    'hard_seat': parts[29] if len(parts) > 29 else '',      # 硬座
                    'no_seat': parts[26] if len(parts) > 26 else '',        # 无座
                }
                
                # 过滤掉无效数据
                if train['train_number'] and train['departure_time']:
                    trains.append(train)
                    
            except Exception as e:
                continue
        
        return trains
    
    def get_train_schedule(self, train_number, date=None):
        """获取列车时刻表"""
        if date is None:
            date = datetime.now().strftime('%Y-%m-%d')
        
        try:
            url = f"{self.base_url}/czxx/query"
            params = {
                'train_no': train_number,
                'from_station_telecode': '',
                'to_station_telecode': '',
                'depart_date': date
            }
            
            headers = {
                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
                'Referer': 'https://kyfw.12306.cn/otn/leftTicket/init'
            }
            
            response = self.session.get(url, params=params, headers=headers, timeout=10)
            
            if response.status_code == 200:
                data = response.json()
                return {
                    'success': True,
                    'schedule': data.get('data', [])
                }
            else:
                return {
                    'success': False,
                    'error': f'请求失败: {response.status_code}'
                }
                
        except Exception as e:
            return {
                'success': False,
                'error': f'查询异常: {str(e)}'
            }

# 测试代码
if __name__ == "__main__":
    railway = RailwayAPI()
    
    # 测试搜索火车票
    result = railway.search_trains('大连', '合肥', '2025-09-16')
    print("搜索结果:")
    print(json.dumps(result, ensure_ascii=False, indent=2))
    
    if result['success'] and result['trains']:
        print(f"\n找到 {len(result['trains'])} 趟列车:")
        for train in result['trains'][:5]:  # 显示前5趟
            print(f"车次: {train['train_number']}")
            print(f"出发时间: {train['departure_time']}")
            print(f"到达时间: {train['arrival_time']}")
            print(f"历时: {train['duration']}")
            print(f"二等座: {train['second_class']}")
            print(f"一等座: {train['first_class']}")
            print("---")